#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
YOLO数据集获取工具 - 主程序入口
基于PyQt5的GUI应用程序，支持数据集网站访问和图片搜索功能
"""

import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QSplitter, QListWidget, QListWidgetItem,
                             QPushButton, QLabel, QTextEdit, QLineEdit, QFrame,
                             QScrollArea, QGridLayout, QMessageBox, QFileDialog,
                             QProgressBar, QStatusBar, QMenuBar, QAction, QToolBar)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer, QUrl
from PyQt5.QtWebEngineWidgets import QWebEngineView
from PyQt5.QtGui import QIcon, QFont, QPixmap, QPalette, QColor
import webbrowser
import json
import requests
from urllib.parse import urljoin, urlparse
import time
from image_search import ImageSearchWorker
from waterfall_display_stable import StableWaterfallDisplay as WaterfallDisplayComponent

class BatchDownloadWorker(QThread):
    """批量下载工作线程"""
    progress_updated = pyqtSignal(int, int, str)  # 当前, 总数, 文件名
    download_finished = pyqtSignal(int, int)  # 成功数, 总数
    
    def __init__(self, images, save_dir):
        super().__init__()
        self.images = images
        self.save_dir = save_dir
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36',
            'Referer': 'https://image.baidu.com',
            'Accept': 'image/webp,image/apng,image/*,*/*;q=0.8',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
            'Connection': 'keep-alive'
        }
        
    def run(self):
        """执行批量下载"""
        success_count = 0
        total_count = len(self.images)
        
        print(f"\n📥 [批量下载] 开始下载 {total_count} 张图片到 {self.save_dir}")
        
        for i, image in enumerate(self.images):
            # 检查线程是否被请求停止
            if self.isInterruptionRequested():
                print("🛑 [BatchDownloadWorker] 线程被请求停止")
                return
                
            try:
                # 生成文件名
                filename = f"{i+1}_{self._sanitize_filename(image.get('title', 'image'))}.jpg"
                save_path = os.path.join(self.save_dir, filename)
                
                print(f"📥 [批量下载] 正在下载 {i+1}/{total_count}: {filename}")
                
                # 下载图片
                url = image.get('thumbnail', image.get('url', ''))
                if url:
                    print(f"🌐 [批量下载] 请求URL: {url[:50]}...")
                    response = requests.get(url, headers=self.headers, timeout=10)
                    print(f"📊 [批量下载] 响应状态码: {response.status_code}, 内容长度: {len(response.content)} bytes")
                    
                    response.raise_for_status()
                    
                    with open(save_path, 'wb') as f:
                        f.write(response.content)
                    
                    success_count += 1
                    print(f"✅ [批量下载] 成功下载: {filename}")
                    self.progress_updated.emit(i+1, total_count, filename)
                    
                    # 避免请求过于频繁
                    time.sleep(0.5)
                else:
                    print(f"❌ [批量下载] 没有有效的图片URL")
                    self.progress_updated.emit(i+1, total_count, f"失败: {filename}")
                    
            except Exception as e:
                print(f"❌ [批量下载] 下载图片失败: {e}")
                print(f"🔍 [批量下载] 错误类型: {type(e).__name__}")
                self.progress_updated.emit(i+1, total_count, f"失败: {filename}")
                
        print(f"🎯 [批量下载] 下载完成: {success_count}/{total_count} 张图片")
        self.download_finished.emit(success_count, total_count)
        
    def _sanitize_filename(self, filename):
        """清理文件名，移除非法字符"""
        import re
        # 移除或替换非法字符
        filename = re.sub(r'[<>:"/\\|?*]', '_', filename)
        # 限制长度
        if len(filename) > 50:
            filename = filename[:50]
        return filename


# WaterfallDisplay类已移至waterfall_display.py模块

class MainWindow(QMainWindow):
    """主窗口"""
    
    def __init__(self):
        super().__init__()
        self.setWindowTitle("YOLO数据集获取工具")
        self.setGeometry(100, 100, 1200, 800)
        self.setup_ui()
        self.setup_menus()
        self.setup_status_bar()
        
    def setup_ui(self):
        """设置UI"""
        # 创建中央widget
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建主布局
        main_layout = QHBoxLayout(central_widget)
        
        # 创建分割器
        splitter = QSplitter(Qt.Horizontal)
        main_layout.addWidget(splitter)
        
        # 左侧菜单
        self.setup_left_sidebar(splitter)
        
        # 中间浏览器区域
        self.setup_center_area(splitter)
        
        # 设置分割器比例
        splitter.setSizes([300, 900])
        
    def setup_left_sidebar(self, parent):
        """设置左侧边栏"""
        sidebar_widget = QWidget()
        sidebar_layout = QVBoxLayout(sidebar_widget)
        
        # 数据集网站分类
        dataset_label = QLabel("📊 数据集网站")
        dataset_label.setFont(QFont("Arial", 12, QFont.Bold))
        sidebar_layout.addWidget(dataset_label)
        
        self.dataset_list = QListWidget()
        self.dataset_list.setMaximumHeight(200)
        dataset_sites = [
            "COCO", "ImageNet", "Open Images", "Pascal VOC",
            "YOLO Dataset", "Roboflow", "Kaggle", "Papers With Code"
        ]
        for site in dataset_sites:
            item = QListWidgetItem(site)
            self.dataset_list.addItem(item)
        self.dataset_list.itemClicked.connect(self.on_dataset_clicked)
        sidebar_layout.addWidget(self.dataset_list)
        
        # 图片搜索网站分类
        search_label = QLabel("🔍 图片搜索")
        search_label.setFont(QFont("Arial", 12, QFont.Bold))
        sidebar_layout.addWidget(search_label)
        
        self.search_list = QListWidget()
        search_sites = [
            "百度图片", "Bing Images"
        ]
        for site in search_sites:
            item = QListWidgetItem(site)
            self.search_list.addItem(item)
        self.search_list.itemClicked.connect(self.on_search_clicked)
        sidebar_layout.addWidget(self.search_list)
        
        parent.addWidget(sidebar_widget)
        
    def setup_center_area(self, parent):
        """设置中间区域"""
        center_widget = QWidget()
        center_layout = QVBoxLayout(center_widget)
        
        # 搜索框（图片搜索时显示）
        self.search_frame = QFrame()
        self.search_frame.setVisible(False)
        search_layout = QHBoxLayout(self.search_frame)
        
        self.search_input = QLineEdit()
        self.search_input.setPlaceholderText("输入搜索关键词...")
        search_layout.addWidget(self.search_input)
        
        self.search_button = QPushButton("搜索")
        self.search_button.clicked.connect(self.perform_search)
        search_layout.addWidget(self.search_button)
        
        center_layout.addWidget(self.search_frame)
        
        # 浏览器/瀑布流显示区域
        self.content_stack = QWidget()
        self.content_layout = QVBoxLayout(self.content_stack)
        
        # Web浏览器
        self.web_view = QWebEngineView()
        self.content_layout.addWidget(self.web_view)
        
        # 瀑布流显示（图片搜索时使用）
        self.waterfall_display = WaterfallDisplayComponent()
        self.waterfall_display.setVisible(False)
        self.content_layout.addWidget(self.waterfall_display)
        
        center_layout.addWidget(self.content_stack)
        
        # 下载控制区域
        self.setup_download_area(center_layout)
        
        parent.addWidget(center_widget)
        
    def setup_download_area(self, parent_layout):
        """设置下载控制区域"""
        download_frame = QFrame()
        download_frame.setFrameStyle(QFrame.StyledPanel)
        download_layout = QHBoxLayout(download_frame)
        
        # 选中图片信息
        self.selected_info = QLabel("已选中 0 张图片")
        download_layout.addWidget(self.selected_info)
        
        download_layout.addStretch()
        
        # 分页控件
        self.setup_pagination_controls(download_layout)
        
        download_layout.addStretch()
        
        # 下载按钮
        self.download_button = QPushButton("下载选中图片")
        self.download_button.clicked.connect(self.download_selected_images)
        self.download_button.setEnabled(False)
        download_layout.addWidget(self.download_button)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        download_layout.addWidget(self.progress_bar)
        
        parent_layout.addWidget(download_frame)
        
    def setup_pagination_controls(self, parent_layout):
        """设置分页控件"""
        # 分页信息
        self.page_info = QLabel("第 1 页")
        parent_layout.addWidget(self.page_info)
        
        # 上一页按钮
        self.prev_button = QPushButton("上一页")
        self.prev_button.clicked.connect(self.prev_page)
        self.prev_button.setEnabled(False)
        parent_layout.addWidget(self.prev_button)
        
        # 下一页按钮
        self.next_button = QPushButton("下一页")
        self.next_button.clicked.connect(self.next_page)
        self.next_button.setEnabled(False)
        parent_layout.addWidget(self.next_button)
        
        # 跳转页面输入框
        self.page_input = QLineEdit()
        self.page_input.setPlaceholderText("页码")
        self.page_input.setMaximumWidth(60)
        self.page_input.returnPressed.connect(self.jump_to_page)
        parent_layout.addWidget(self.page_input)
        
        # 跳转按钮
        self.jump_button = QPushButton("跳转")
        self.jump_button.clicked.connect(self.jump_to_page)
        parent_layout.addWidget(self.jump_button)
        
        # 初始化分页状态
        self.current_page = 1
        self.total_pages = 0
        self.has_more_pages = True
        
    def setup_menus(self):
        """设置菜单栏"""
        menubar = self.menuBar()
        
        # 文件菜单
        file_menu = menubar.addMenu('文件')
        
        new_action = QAction('新建标签页', self)
        new_action.setShortcut('Ctrl+N')
        new_action.triggered.connect(self.new_tab)
        file_menu.addAction(new_action)
        
        # 编辑菜单
        edit_menu = menubar.addMenu('编辑')
        
        refresh_action = QAction('刷新页面', self)
        refresh_action.setShortcut('Ctrl+R')
        refresh_action.triggered.connect(self.refresh_page)
        edit_menu.addAction(refresh_action)
        
        # 工具菜单
        tools_menu = menubar.addMenu('工具')
        
        download_action = QAction('下载文件', self)
        download_action.setShortcut('Ctrl+D')
        download_action.triggered.connect(self.download_selected_images)
        tools_menu.addAction(download_action)
        
    def setup_status_bar(self):
        """设置状态栏"""
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        self.status_bar.showMessage("就绪")
        
    def on_dataset_clicked(self, item):
        """数据集网站点击事件"""
        site_name = item.text()
        urls = {
            "COCO": "https://cocodataset.org/",
            "ImageNet": "https://www.image-net.org/",
            "Open Images": "https://storage.googleapis.com/openimages/web/index.html",
            "Pascal VOC": "http://host.robots.ox.ac.uk/pascal/VOC/",
            "YOLO Dataset": "https://github.com/ultralytics/yolov5",
            "Roboflow": "https://roboflow.com/",
            "Kaggle": "https://www.kaggle.com/datasets",
            "Papers With Code": "https://paperswithcode.com/datasets"
        }
        
        if site_name in urls:
            # 使用外部浏览器打开
            webbrowser.open(urls[site_name])
            self.status_bar.showMessage(f"已在外部浏览器中打开 {site_name}")
            
    def on_search_clicked(self, item):
        """图片搜索网站点击事件"""
        site_name = item.text()
        
        # 显示搜索框
        self.search_frame.setVisible(True)
        
        # 隐藏Web浏览器，显示瀑布流
        self.web_view.setVisible(False)
        self.waterfall_display.setVisible(True)
        
        # 清空之前的搜索结果
        self.waterfall_display.clear_images()
        
        self.status_bar.showMessage(f"已切换到 {site_name} 搜索模式")
        
    def perform_search(self):
        """执行图片搜索"""
        search_term = self.search_input.text().strip()
        if not search_term:
            QMessageBox.warning(self, "警告", "请输入搜索关键词")
            return
            
        # 获取当前选中的搜索引擎
        current_item = self.search_list.currentItem()
        if not current_item:
            QMessageBox.warning(self, "警告", "请先选择搜索引擎")
            return
            
        search_engine = current_item.text()
        
        print(f"\n🚀 [主程序] 开始图片搜索")
        print(f"🔍 [主程序] 搜索引擎: {search_engine}")
        print(f"📝 [主程序] 搜索关键词: {search_term}")
        
        # 重置分页状态
        self.current_page = 1
        self.total_pages = 0
        self.has_more_pages = True
        
        # 显示进度条
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        
        # 启动搜索线程（只搜索第1页，每页10张图片）
        self.search_worker = ImageSearchWorker(search_term, search_engine, max_pages=1, per_page=10, start_page=1)
        self.search_worker.search_finished.connect(self.on_search_finished)
        self.search_worker.progress_updated.connect(self.progress_bar.setValue)
        self.search_worker.start()
        
        self.status_bar.showMessage(f"正在使用 {search_engine} 搜索: {search_term}")
        
    def on_search_finished(self, images, error):
        """搜索完成回调"""
        self.progress_bar.setVisible(False)
        
        print(f"\n🎯 [主程序] 搜索完成回调")
        print(f"📊 [主程序] 搜索结果数量: {len(images) if images else 0}")
        print(f"❌ [主程序] 错误信息: {error if error else '无'}")
        
        if error:
            print(f"❌ [主程序] 搜索失败: {error}")
            QMessageBox.critical(self, "错误", f"搜索失败: {error}")
            return
            
        # 更新分页状态
        if images and len(images) > 0:
            # 如果返回的图片数量等于请求的数量，可能还有更多页
            self.has_more_pages = len(images) >= 10
            if self.has_more_pages:
                self.total_pages = self.current_page + 1  # 至少还有下一页
            else:
                self.total_pages = self.current_page
                self.has_more_pages = False
        else:
            # 没有图片，说明没有更多页了
            self.has_more_pages = False
            self.total_pages = max(1, self.current_page - 1)
            
        # 更新分页控件状态
        self.update_pagination_controls()
            
        # 添加图片到瀑布流
        print(f"✅ [主程序] 开始添加 {len(images)} 张图片到瀑布流")
        self.waterfall_display.add_images(images)
        
        # 显示瀑布流组件
        self.waterfall_display.setVisible(True)
        print(f"👁️ [主程序] 瀑布流组件已设置为可见")
        
        # 连接信号
        self.waterfall_display.selection_changed.connect(self.on_selection_changed)
        
        print(f"🎯 [主程序] 搜索完成，找到 {len(images)} 张图片")
        self.status_bar.showMessage(f"搜索完成，找到 {len(images)} 张图片")
        
    def update_pagination_controls(self):
        """更新分页控件状态"""
        # 更新页面信息
        if self.has_more_pages:
            self.page_info.setText(f"第 {self.current_page} 页")
        else:
            self.page_info.setText(f"第 {self.current_page} 页 (图片已经全部显示)")
            
        # 更新按钮状态
        self.prev_button.setEnabled(self.current_page > 1)
        self.next_button.setEnabled(self.has_more_pages)
        
        # 清空跳转输入框
        self.page_input.clear()
        
    def prev_page(self):
        """上一页"""
        if self.current_page > 1:
            self.current_page -= 1
            self.load_page()
            
    def next_page(self):
        """下一页"""
        if self.has_more_pages:
            self.current_page += 1
            self.load_page()
            
    def jump_to_page(self):
        """跳转到指定页面"""
        try:
            page_num = int(self.page_input.text())
            if page_num >= 1 and page_num != self.current_page:
                self.current_page = page_num
                self.load_page()
        except ValueError:
            QMessageBox.warning(self, "警告", "请输入有效的页码")
            
    def load_page(self):
        """加载指定页面"""
        search_term = self.search_input.text().strip()
        if not search_term:
            return
            
        # 获取当前选中的搜索引擎
        current_item = self.search_list.currentItem()
        if not current_item:
            return
            
        search_engine = current_item.text()
        
        print(f"\n📄 [分页] 加载第 {self.current_page} 页")
        
        # 显示进度条
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        
        # 清空当前图片
        self.waterfall_display.clear_images()
        
        # 启动搜索线程（只搜索当前页，每页10张图片）
        self.search_worker = ImageSearchWorker(search_term, search_engine, max_pages=1, per_page=10, start_page=self.current_page)
        self.search_worker.search_finished.connect(self.on_page_loaded)
        self.search_worker.progress_updated.connect(self.progress_bar.setValue)
        self.search_worker.start()
        
        self.status_bar.showMessage(f"正在加载第 {self.current_page} 页...")
        
    def on_page_loaded(self, images, error):
        """页面加载完成回调"""
        self.progress_bar.setVisible(False)
        
        if error:
            print(f"❌ [分页] 页面加载失败: {error}")
            QMessageBox.critical(self, "错误", f"页面加载失败: {error}")
            return
            
        # 更新分页状态
        if images and len(images) > 0:
            # 如果返回的图片数量等于请求的数量，可能还有更多页
            self.has_more_pages = len(images) >= 10
            if self.has_more_pages:
                self.total_pages = self.current_page + 1  # 至少还有下一页
            else:
                self.total_pages = self.current_page
                self.has_more_pages = False
        else:
            # 没有图片，说明没有更多页了
            self.has_more_pages = False
            self.total_pages = max(1, self.current_page - 1)
            
        # 更新分页控件状态
        self.update_pagination_controls()
            
        # 添加图片到瀑布流
        print(f"✅ [分页] 添加 {len(images)} 张图片到瀑布流")
        self.waterfall_display.add_images(images)
        
        # 连接信号
        self.waterfall_display.selection_changed.connect(self.on_selection_changed)
        
        print(f"🎯 [分页] 第 {self.current_page} 页加载完成，找到 {len(images)} 张图片")
        self.status_bar.showMessage(f"第 {self.current_page} 页加载完成，找到 {len(images)} 张图片")
        
    def on_selection_changed(self, selected_images):
        """选择变化事件"""
        selected_count = len(selected_images)
        self.selected_info.setText(f"已选中 {selected_count} 张图片")
        self.download_button.setEnabled(selected_count > 0)
        
    def download_selected_images(self):
        """下载选中的图片"""
        selected_images = self.waterfall_display.get_selected_images()
        if not selected_images:
            QMessageBox.warning(self, "警告", "请先选择要下载的图片")
            return
            
        # 选择保存目录
        save_dir = QFileDialog.getExistingDirectory(self, "选择保存目录")
        if not save_dir:
            return
            
        # 创建下载目录
        search_term = self.search_input.text().strip() or "images"
        download_dir = os.path.join(save_dir, search_term)
        os.makedirs(download_dir, exist_ok=True)
        
        # 启动批量下载
        self.batch_download_images(selected_images, download_dir)
        
    def batch_download_images(self, images, save_dir):
        """批量下载图片"""
        self.download_progress = QProgressBar()
        self.download_progress.setVisible(True)
        self.download_progress.setMaximum(len(images))
        self.download_progress.setValue(0)
        
        # 添加到下载控制区域
        self.content_layout.addWidget(self.download_progress)
        
        # 创建下载工作线程
        self.batch_download_worker = BatchDownloadWorker(images, save_dir)
        self.batch_download_worker.progress_updated.connect(self.on_batch_download_progress)
        self.batch_download_worker.download_finished.connect(self.on_batch_download_finished)
        self.batch_download_worker.start()
        
    def on_batch_download_progress(self, current, total, filename):
        """批量下载进度更新"""
        self.download_progress.setValue(current)
        self.status_bar.showMessage(f"正在下载 {filename} ({current}/{total})")
        
    def on_batch_download_finished(self, success_count, total_count):
        """批量下载完成"""
        self.download_progress.setVisible(False)
        self.status_bar.showMessage(f"下载完成: {success_count}/{total_count} 张图片")
        QMessageBox.information(self, "下载完成", f"成功下载 {success_count} 张图片，共 {total_count} 张")
        
    def new_tab(self):
        """新建标签页"""
        self.web_view.setUrl(QUrl("about:blank"))
        self.status_bar.showMessage("新建标签页")
        
    def refresh_page(self):
        """刷新页面"""
        if self.web_view.isVisible():
            self.web_view.reload()
        self.status_bar.showMessage("页面已刷新")
    
    def closeEvent(self, event):
        """窗口关闭事件，清理资源"""
        print("🔄 [主程序] 开始清理资源...")
        
        # 停止搜索工作线程
        if hasattr(self, 'search_worker') and self.search_worker.isRunning():
            print("🛑 [主程序] 停止搜索工作线程")
            self.search_worker.requestInterruption()
            self.search_worker.wait(3000)  # 等待3秒
            
        # 停止批量下载工作线程
        if hasattr(self, 'batch_download_worker') and self.batch_download_worker.isRunning():
            print("🛑 [主程序] 停止批量下载工作线程")
            self.batch_download_worker.requestInterruption()
            self.batch_download_worker.wait(3000)  # 等待3秒
            
        # 清理瀑布流组件
        if hasattr(self, 'waterfall_display'):
            print("🧹 [主程序] 清理瀑布流组件")
            self.waterfall_display.clear_images()
            
        # 清理WebEngineView
        if hasattr(self, 'web_view'):
            print("🧹 [主程序] 清理WebEngineView")
            self.web_view.setUrl(QUrl("about:blank"))
            
        print("✅ [主程序] 资源清理完成")
        event.accept()

def main():
    """主函数"""
    app = QApplication(sys.argv)
    
    # 设置应用程序信息
    app.setApplicationName("YOLO数据集获取工具")
    app.setApplicationVersion("1.0.0")
    app.setOrganizationName("YOLO Dataset Tool")
    
    # 创建主窗口
    window = MainWindow()
    window.show()
    
    # 运行应用程序
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()
